# Basic Functions by Tao YiNing
import cv2
import numpy as np

#####################定位#####################
# 预处理
def preprocess(img1, img2):
    # 彩色图->灰度图
    im1 = cv2.cvtColor(img1, cv2.COLOR_BGR2GRAY)
    im2 = cv2.cvtColor(img2, cv2.COLOR_BGR2GRAY)
    # 直方图均衡化
    im1 = cv2.equalizeHist(im1)
    im2 = cv2.equalizeHist(im2)
    return im1, im2 

# 消除畸变
# 获取畸变校正和立体校正的映射变换矩阵、重投影矩阵
# stereo_config是一个类，存储着双目标定的参数:config = stereo_config.stereoCamera()
def getRectifyTransform(height, width, config):
    # 读取内参和外参
    global left_K,left_distortion,R1,P1
    global right_K, right_distortion
    left_K = config.cam_matrix_left
    right_K = config.cam_matrix_right
    left_distortion = config.distortion_l
    right_distortion = config.distortion_r
    R = config.R
    T = config.T
    
    # 计算校正变换
    # R1为旋转矩阵，相机坐标系的校正；P1为新投影矩阵，即新相机坐标系点投影于图像坐标系；
    R1, R2, P1, P2, Q, roi1, roi2 = cv2.stereoRectify(left_K, left_distortion, right_K, right_distortion, (width, height), R, T, alpha=0)
    map1x, map1y = cv2.initUndistortRectifyMap(left_K, left_distortion, R1, P1, (width, height), cv2.CV_32FC1)
    map2x, map2y = cv2.initUndistortRectifyMap(right_K, right_distortion, R2, P2, (width, height), cv2.CV_32FC1)
    return map1x, map1y, map2x, map2y, Q 

# 畸变校正和立体校正
def rectifyImage(image1, image2, map1x, map1y, map2x, map2y):
    rectifyed_img1 = cv2.remap(image1, map1x, map1y, cv2.INTER_AREA)
    rectifyed_img2 = cv2.remap(image2, map2x, map2y, cv2.INTER_AREA)
    return rectifyed_img1, rectifyed_img2

# 图像识别框坐标校正
def rectifyboxes(detecting_box):
    src = np.zeros(shape = (2,1,2))
    src[0][0][0] = detecting_box[0]
    src[0][0][1] = detecting_box[1]
    src[1][0][0] = detecting_box[2]
    src[1][0][1] = detecting_box[3]
    
    points_reflicted = cv2.undistortPoints(src, left_K, left_distortion,R=R1,P=P1)
    return points_reflicted

# added by fushuai
def undistortPoints(pix_x, pix_y, l_or_r):
    src = np.zeros(shape = (1,1,2))
    src[0][0][0] = pix_x
    src[0][0][1] = pix_y
    if l_or_r == "left":
        [[[x_out, y_out]]] = cv2.undistortPoints(src, left_K, left_distortion, R=R1, P=P1)
    elif l_or_r == "right":
        [[[x_out, y_out]]] = cv2.undistortPoints(src, right_K, right_distortion, R=R1, P=P1)
    return x_out, y_out

#定位时选取五个点以保证中心点不为空，减小误差
def five_points(box_rectify):
    left_up_X = box_rectify[0, 0, 0]
    left_up_Y = box_rectify[0, 0, 1]
    right_down_X = box_rectify[1, 0, 0]
    right_down_Y = box_rectify[1, 0, 1]
    w = right_down_X - left_up_X  #识别框宽度
    h = right_down_Y - left_up_Y  #识别框高度
    uppoint = [int(left_up_X + w /2), int(left_up_Y + h /4)]
    downpoint = [int(left_up_X + w /2), int(left_up_Y + 3*h /4)]
    centerpoint = [int(left_up_X + w /2), int(left_up_Y + h /2)]
    leftpoint = [int(left_up_X + w /4), int(left_up_Y + h /2)]
    rightpoint = [int(left_up_X + 3*w /4), int(left_up_Y + h /2)]
    
    point =[uppoint, downpoint, rightpoint, leftpoint, centerpoint]
    return point
    
# 视差计算SGBM
def disparity_SGBM(left_image, right_image, down_scale=False):
    # SGBM匹配参数设置
    if left_image.ndim == 2:
        img_channels = 1
    else:
        img_channels = 3
    blockSize = 3
    param = {'minDisparity': 0,  #最小视差，默认0
             'numDisparities': 128,  #视差搜索范围，16的整数倍
             'blockSize': blockSize,  #块匹配的大小。奇数，3-11
             'P1': 8 * img_channels * blockSize ** 2,  #值越大，视差越平滑
             'P2': 32 * img_channels * blockSize ** 2,  #值越大，视差越平滑
             'disp12MaxDiff': 1,  #左右一致性检测中最大容许误差值
             'preFilterCap': 63,  #映射滤波器大小，默认15
             'uniquenessRatio': 15,  #唯一性检测参数，匹配度分区不够则误匹配
             'speckleWindowSize': 100,  #小于此则为噪声点
             'speckleRange': 1,  #大于此则为不连通
             'mode': cv2.STEREO_SGBM_MODE_SGBM_3WAY
             }
    # 构建SGBM对象
    sgbm = cv2.StereoSGBM_create(**param)
    # 计算视差图
    size = (left_image.shape[1], left_image.shape[0])
    if down_scale == False:
        disparity_left = sgbm.compute(left_image, right_image)
        disparity_right = sgbm.compute(right_image, left_image)
    else:
        left_image_down = cv2.pyrDown(left_image)
        right_image_down = cv2.pyrDown(right_image)
        factor = size[0] / left_image_down.shape[1]
        disparity_left_half = sgbm.compute(left_image_down, right_image_down)
        disparity_right_half = sgbm.compute(right_image_down, left_image_down)
        disparity_left = cv2.resize(disparity_left_half, size, interpolation=cv2.INTER_AREA) 
        disparity_right = cv2.resize(disparity_right_half, size, interpolation=cv2.INTER_AREA)
        disparity_left *= factor 
        disparity_right *= factor
    return disparity_left, disparity_right